# api/task_manager.py

import os
import shutil
from datetime import datetime, timedelta
from typing import Dict, Optional
from uuid import uuid4
import asyncio
import logging


class TaskManager:
    """
    异步任务管理器
    
    单例模式，管理异步任务的创建、状态跟踪、并发控制和自动清理
    """
    
    _instance = None
    _lock = asyncio.Lock()
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(TaskManager, cls).__new__(cls)
            cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if not self._initialized:
            self.tasks: Dict[str, Dict] = {}
            self.max_concurrent_tasks = 3
            self.task_retention_hours = 24
            self.temp_base_path = "/root/FaceFusion/temp"
            self._cleanup_lock = asyncio.Lock()
            self._initialized = True
            
            # 确保临时目录存在
            os.makedirs(self.temp_base_path, exist_ok=True)
            
            # 配置日志
            logging.basicConfig(
                level=logging.INFO,
                format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            self.logger = logging.getLogger("TaskManager")
    
    async def create_task(self, source_path: str, target_path: str, mode: str, output_quality: int) -> str:
        """
        创建新任务
        
        Returns:
            str: 任务ID
            
        Raises:
            ValueError: 参数无效时
        """
        if not source_path or not target_path:
            raise ValueError("source_path and target_path are required")
        
        if output_quality < 1 or output_quality > 100:
            raise ValueError("output_quality must be between 1 and 100")
        
        async with self._lock:
            task_id = str(uuid4())
            task_temp_path = os.path.join(self.temp_base_path, task_id)
            
            # 创建任务临时目录
            os.makedirs(task_temp_path, exist_ok=True)
            
            task_data = {
                "task_id": task_id,
                "status": "pending",
                "progress": 0,
                "source_path": source_path,
                "target_path": target_path,
                "result_path": None,
                "error": None,
                "created_at": datetime.now(),
                "updated_at": datetime.now(),
                "mode": mode,
                "output_quality": output_quality,
                "temp_path": task_temp_path
            }
            
            self.tasks[task_id] = task_data
            
            self.logger.info(f"Task created: {task_id}, mode: {mode}, quality: {output_quality}")
            
            return task_id
    
    async def get_task(self, task_id: str) -> Dict:
        """
        获取任务信息
        
        Args:
            task_id: 任务ID
            
        Returns:
            Dict: 任务数据
            
        Raises:
            ValueError: 任务不存在时
        """
        async with self._lock:
            if task_id not in self.tasks:
                raise ValueError(f"Task not found: {task_id}")
            
            return self.tasks[task_id].copy()
    
    async def update_task_status(self, task_id: str, status: str, progress: Optional[int] = None, 
                               error: Optional[str] = None) -> None:
        """
        更新任务状态
        
        Args:
            task_id: 任务ID
            status: 新状态 (pending/processing/completed/failed)
            progress: 进度值 (0-100)
            error: 错误信息
            
        Raises:
            ValueError: 任务不存在或状态无效时
        """
        valid_statuses = {"pending", "processing", "completed", "failed"}
        if status not in valid_statuses:
            raise ValueError(f"Invalid status: {status}. Must be one of {valid_statuses}")
        
        async with self._lock:
            if task_id not in self.tasks:
                raise ValueError(f"Task not found: {task_id}")
            
            task = self.tasks[task_id]
            task["status"] = status
            task["updated_at"] = datetime.now()
            
            if progress is not None:
                if progress < 0 or progress > 100:
                    raise ValueError("Progress must be between 0 and 100")
                task["progress"] = progress
            
            if error is not None:
                task["error"] = error
            
            self.logger.info(f"Task {task_id} status updated to {status}, progress: {task['progress']}")
            
            # 如果任务失败，清理临时文件
            if status == "failed":
                await self._cleanup_task_files(task_id)
    
    async def set_task_result(self, task_id: str, result_path: str) -> None:
        """
        设置任务结果路径
        
        Args:
            task_id: 任务ID
            result_path: 结果文件路径
            
        Raises:
            ValueError: 任务不存在时
        """
        async with self._lock:
            if task_id not in self.tasks:
                raise ValueError(f"Task not found: {task_id}")
            
            self.tasks[task_id]["result_path"] = result_path
            self.tasks[task_id]["updated_at"] = datetime.now()
            
            self.logger.info(f"Task {task_id} result set: {result_path}")
    
    async def cleanup_old_tasks(self) -> None:
        """
        清理过期任务（24小时前创建的任务）
        """
        async with self._cleanup_lock:
            cutoff_time = datetime.now() - timedelta(hours=self.task_retention_hours)
            tasks_to_remove = []
            
            async with self._lock:
                for task_id, task in self.tasks.items():
                    if task["created_at"] < cutoff_time:
                        tasks_to_remove.append(task_id)
                
                for task_id in tasks_to_remove:
                    await self._cleanup_task_files(task_id)
                    del self.tasks[task_id]
            
            if tasks_to_remove:
                self.logger.info(f"Cleaned up {len(tasks_to_remove)} old tasks")
    
    async def get_active_tasks_count(self) -> int:
        """
        获取正在处理的任务数量
        
        Returns:
            int: 正在处理的任务数量
        """
        async with self._lock:
            return sum(1 for task in self.tasks.values() if task["status"] == "processing")
    
    async def get_pending_tasks_count(self) -> int:
        """
        获取待处理的任务数量
        
        Returns:
            int: 待处理的任务数量
        """
        async with self._lock:
            return sum(1 for task in self.tasks.values() if task["status"] == "pending")
    
    async def can_process_new_task(self) -> bool:
        """
        检查是否可以处理新任务（并发控制）
        
        Returns:
            bool: 是否可以处理新任务
        """
        active_count = await self.get_active_tasks_count()
        return active_count < self.max_concurrent_tasks
    
    async def _cleanup_task_files(self, task_id: str) -> None:
        """
        清理任务临时文件
        
        Args:
            task_id: 任务ID
        """
        try:
            task = self.tasks.get(task_id)
            if task and "temp_path" in task:
                temp_path = task["temp_path"]
                if os.path.exists(temp_path):
                    shutil.rmtree(temp_path)
                    self.logger.info(f"Cleaned up temp files for task: {task_id}")
        except Exception as e:
            self.logger.error(f"Error cleaning up task files for {task_id}: {str(e)}")
    
    async def get_all_tasks(self) -> Dict[str, Dict]:
        """
        获取所有任务（用于调试和管理）
        
        Returns:
            Dict[str, Dict]: 所有任务的字典
        """
        async with self._lock:
            return {task_id: task_data.copy() for task_id, task_data in self.tasks.items()}
    
    async def complete_task(self, task_id: str, result_path: str) -> None:
        """
        标记任务为完成状态
        
        Args:
            task_id: 任务ID
            result_path: 结果文件路径
        """
        await self.update_task_status(task_id, "completed", progress=100)
        await self.set_task_result(task_id, result_path)
        self.logger.info(f"Task completed: {task_id}")
    
    async def fail_task(self, task_id: str, error_message: str) -> None:
        """
        标记任务为失败状态
        
        Args:
            task_id: 任务ID
            error_message: 错误信息
        """
        await self.update_task_status(task_id, "failed", error=error_message)
        self.logger.error(f"Task failed: {task_id}, error: {error_message}")
    
    async def cancel_task(self, task_id: str) -> bool:
        """
        取消任务
        
        Args:
            task_id: 任务ID
            
        Returns:
            bool: 是否成功取消
            
        Raises:
            ValueError: 任务不存在时
        """
        async with self._lock:
            if task_id not in self.tasks:
                raise ValueError(f"Task not found: {task_id}")
            
            task = self.tasks[task_id]
            if task["status"] in {"pending", "processing"}:
                task["status"] = "failed"
                task["error"] = "Task cancelled by user"
                task["updated_at"] = datetime.now()
                
                await self._cleanup_task_files(task_id)
                self.logger.info(f"Task cancelled: {task_id}")
                return True
            
            return False
